
#include <linux/module.h>
#include <linux/irq.h>
#include <linux/platform_device.h>
#include <linux/dma-mapping.h>
#include <linux/cdev.h>
#include <linux/types.h>
#include <linux/mutex.h>
#include <linux/acpi.h>
#include <crypto/internal/hash.h>
#include <linux/crypto.h>
#include <crypto/hash.h>
#include <crypto/internal/skcipher.h>
#include <crypto/internal/akcipher.h>
#include <linux/ioport.h>
#include <linux/io.h>
#include <linux/oid_registry.h>

#include "test.h"

#include "sm2_test.h"

static u8 *test_pack_u32(u8 *dst, u32 val)
{
	memcpy(dst, &val, sizeof(val));
	return dst + sizeof(val);
}

static int ecdsa_test_encrypt(char *alg, struct akcipher_testvec *vec)
{
	struct crypto_akcipher *tfm;
	char *src_buffer = NULL;
	char *dst_buffer = NULL;
	struct akcipher_request *req;
	struct crypto_wait wait;
	unsigned int out_len_max, out_len = 0;
	struct scatterlist src, dst;
	const char *op;
	u8 *key, *ptr;
	int err = 0;

	tfm = crypto_alloc_akcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR
		       "alg: akcipher: Failed to load tfm for %s: %ld\n",
		       alg, PTR_ERR(tfm));
		return PTR_ERR(tfm);
	}

	req = akcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "akcipher_request_alloc failed");
		goto fail;
	}

	crypto_init_wait(&wait);

	/*  key | algo | param_len | param */
	key = kmalloc(vec->key_len + sizeof(u32) * 2 + vec->param_len,
		      GFP_KERNEL);
	if (!key) {
		printk(KERN_ERR "alloc key failed");

		goto free_req;
	}

	memcpy(key, vec->key, vec->key_len);
	ptr = key + vec->key_len;
	ptr = test_pack_u32(ptr, vec->algo);
	ptr = test_pack_u32(ptr, vec->param_len);
	memcpy(ptr, vec->params, vec->param_len);

	if (!vec->public_key_vec) {
		printk(KERN_ERR "%s line %d skip sign without public key",
		       __FUNCTION__, __LINE__);
		goto free_key;
	}

	err = crypto_akcipher_set_pub_key(tfm, key, vec->key_len);
	if (err) {
		printk(KERN_ERR "crypto_akcipher_set_pub_key key failed");
		goto free_key;
	}

	out_len_max = 1024;
	dst_buffer = kzalloc(out_len_max, GFP_KERNEL);
	if (!dst_buffer) {
		err = -ENOMEM;
		goto free_all;
	}

	err = -E2BIG;
	op = "encrypt";

	if (!vec->m) {
		err = 0;
		printk(KERN_ERR "%s line %d missing digest m", __FUNCTION__,
		       __LINE__);
		goto free_all;
	}

	src_buffer = kzalloc(vec->m_size, GFP_KERNEL);
	if (!src_buffer) {
		goto free_all;
	}

	memcpy(src_buffer, vec->m, vec->m_size);

	sg_init_one(&src, src_buffer, vec->m_size);
	sg_init_one(&dst, dst_buffer, out_len_max);
	crypto_init_wait(&wait);
	akcipher_request_set_crypt(req, &src, &dst, vec->m_size, out_len_max);

	akcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				      crypto_req_done, &wait);

	err = crypto_wait_req(crypto_akcipher_encrypt(req), &wait);
	if (err) {
		printk(KERN_ERR "alg: akcipher: %s test failed. err %d\n", op,
		       err);
		goto free_all;
	}

	out_len = req->dst_len;
	vec->c_size = out_len;
	memcpy(vec->c, dst_buffer, out_len);
	// print_buf("get enc", dst_buffer, out_len);

free_all:
	kfree(src_buffer);
	kfree(dst_buffer);
free_key:
	kfree(key);

free_req:
	akcipher_request_free(req);

fail:
	crypto_free_akcipher(tfm);

	return err;
}

static int ecdsa_test_decrypt(char *alg, struct akcipher_testvec *vec)
{
	struct crypto_akcipher *tfm;
	char *src_buffer = NULL;
	char *dst_buffer = NULL;
	struct akcipher_request *req;
	struct crypto_wait wait;
	unsigned int out_len_max, out_len = 0;
	struct scatterlist src, dst;
	const char *op;
	u8 *key, *ptr;
	int err = 0;

	tfm = crypto_alloc_akcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR
		       "alg: akcipher: Failed to load tfm for %s: %ld\n",
		       alg, PTR_ERR(tfm));
		return PTR_ERR(tfm);
	}

	req = akcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "akcipher_request_alloc failed");
		goto fail;
	}

	crypto_init_wait(&wait);

	/*  key | algo | param_len | param */
	key = kmalloc(vec->key_len + sizeof(u32) * 2 + vec->param_len,
		      GFP_KERNEL);
	if (!key) {
		printk(KERN_ERR "alloc key failed");

		goto free_req;
	}

	memcpy(key, vec->key, vec->key_len);
	ptr = key + vec->key_len;
	ptr = test_pack_u32(ptr, vec->algo);
	ptr = test_pack_u32(ptr, vec->param_len);
	memcpy(ptr, vec->params, vec->param_len);

	if (vec->public_key_vec) {
		printk(KERN_ERR "%s line %d skip sign without private key",
		       __FUNCTION__, __LINE__);
		goto free_key;
	}

	err = crypto_akcipher_set_priv_key(tfm, key, vec->key_len);
	if (err) {
		printk(KERN_ERR "crypto_akcipher_set_priv_key key failed");
		goto free_key;
	}

	out_len_max = 1024;
	dst_buffer = kzalloc(out_len_max, GFP_KERNEL);
	if (!dst_buffer) {
		err = -ENOMEM;
		goto free_all;
	}

	err = -E2BIG;
	op = "decrypt";

	if (!vec->c) {
		err = 0;
		printk(KERN_ERR "%s line %d missing cipher ", __FUNCTION__,
		       __LINE__);
		goto free_all;
	}

	src_buffer = kzalloc(vec->c_size, GFP_KERNEL);
	if (!src_buffer) {
		goto free_all;
	}

	memcpy(src_buffer, vec->c, vec->c_size);

	sg_init_one(&src, src_buffer, vec->c_size);
	sg_init_one(&dst, dst_buffer, out_len_max);
	crypto_init_wait(&wait);
	akcipher_request_set_crypt(req, &src, &dst, vec->c_size, out_len_max);

	akcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				      crypto_req_done, &wait);

	err = crypto_wait_req(crypto_akcipher_decrypt(req), &wait);
	if (err) {
		printk(KERN_ERR "alg: akcipher: %s test failed. err %d\n", op,
		       err);
		goto free_all;
	}

	out_len = req->dst_len;
	if (out_len != vec->m_size) {
		err = -1;
		printk(KERN_ERR
		       "alg: akcipher: %s test failed m_size != out_len. err %d\n",
		       op, err);
		goto free_all;
	}

	if (memcmp(vec->m, dst_buffer, vec->m_size)) {
		err = -1;
		pr_err("alg: akcipher: %s test failed. Invalid output\n", op);
		goto free_all;
	}

free_all:
	kfree(src_buffer);
	kfree(dst_buffer);
free_key:
	kfree(key);

free_req:
	akcipher_request_free(req);

fail:
	crypto_free_akcipher(tfm);

	return err;
}

static int ecdsa_test_sign(char *alg, struct akcipher_testvec *vec)
{
	struct crypto_akcipher *tfm;
	char *src_buffer = NULL;
	char *dst_buffer = NULL;
	struct akcipher_request *req;
	struct crypto_wait wait;
	unsigned int out_len_max, out_len = 0;
	struct scatterlist src, dst;
	const char *op;
	u8 *key, *ptr;
	int err = 0;

	tfm = crypto_alloc_akcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR
		       "alg: akcipher: Failed to load tfm for %s: %ld\n",
		       alg, PTR_ERR(tfm));
		return PTR_ERR(tfm);
	}

	req = akcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "akcipher_request_alloc failed");
		goto fail;
	}

	crypto_init_wait(&wait);

	/*  key | algo | param_len | param */
	key = kmalloc(vec->key_len + sizeof(u32) * 2 + vec->param_len,
		      GFP_KERNEL);
	if (!key) {
		printk(KERN_ERR "alloc key failed");

		goto free_req;
	}

	memcpy(key, vec->key, vec->key_len);
	ptr = key + vec->key_len;
	ptr = test_pack_u32(ptr, vec->algo);
	ptr = test_pack_u32(ptr, vec->param_len);
	memcpy(ptr, vec->params, vec->param_len);

	if (vec->public_key_vec) {
		printk(KERN_ERR "%s line %d skip sign without private key",
		       __FUNCTION__, __LINE__);
		goto free_key;
	}

	err = crypto_akcipher_set_priv_key(tfm, key, vec->key_len);
	if (err) {
		printk(KERN_ERR "crypto_akcipher_set_priv_key key failed");
		goto free_key;
	}

	out_len_max = 1024;
	dst_buffer = kzalloc(out_len_max, GFP_KERNEL);
	if (!dst_buffer) {
		err = -ENOMEM;
		goto free_all;
	}

	err = -E2BIG;
	op = "sign";

	if (!vec->m) {
		err = 0;
		printk(KERN_ERR "%s line %d missing digest m", __FUNCTION__,
		       __LINE__);
		goto free_all;
	}

	src_buffer = kzalloc(vec->m_size, GFP_KERNEL);
	if (!src_buffer) {
		goto free_all;
	}

	memcpy(src_buffer, vec->m, vec->m_size);

	sg_init_one(&src, src_buffer, vec->m_size);
	sg_init_one(&dst, dst_buffer, out_len_max);
	crypto_init_wait(&wait);
	akcipher_request_set_crypt(req, &src, &dst, vec->m_size, out_len_max);

	akcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				      crypto_req_done, &wait);

	err = crypto_wait_req(crypto_akcipher_sign(req), &wait);
	if (err) {
		printk(KERN_ERR "alg: akcipher: %s test failed. err %d\n", op,
		       err);
		goto free_all;
	}

	out_len = req->dst_len;
	vec->c_size = out_len;
	memcpy(vec->c, dst_buffer, out_len);
	// print_buf("get sign", dst_buffer, out_len);

free_all:
	kfree(src_buffer);
	kfree(dst_buffer);
free_key:
	kfree(key);

free_req:
	akcipher_request_free(req);

fail:
	crypto_free_akcipher(tfm);

	return err;
}

static int ecdsa_test_verify(char *alg, struct akcipher_testvec *vec)
{
	struct crypto_akcipher *tfm;
	struct akcipher_request *req;
	struct crypto_wait wait;
	unsigned int out_len_max = 0;
	int err = -ENOMEM;
	char *src_buffer = NULL;
	char *dst_buffer = NULL;

	struct scatterlist src_tab[3];
	u8 *key, *ptr;

	tfm = crypto_alloc_akcipher(alg, 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR
		       "alg: akcipher: Failed to load tfm for %s: %ld\n",
		       alg, PTR_ERR(tfm));
		return PTR_ERR(tfm);
	}

	req = akcipher_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "%s line %d", __FUNCTION__, __LINE__);
		goto free_req;
	}

	crypto_init_wait(&wait);

	key = kmalloc(vec->key_len + sizeof(u32) * 2 + vec->param_len,
		      GFP_KERNEL);
	if (!key) {
		printk(KERN_ERR "%s line %d", __FUNCTION__, __LINE__);

		goto free_req;
	}

	memcpy(key, vec->key, vec->key_len);
	ptr = key + vec->key_len;
	ptr = test_pack_u32(ptr, vec->algo);
	ptr = test_pack_u32(ptr, vec->param_len);
	memcpy(ptr, vec->params, vec->param_len);

	if (!vec->public_key_vec) {
		printk(KERN_ERR "%s line %d skip without public key ",
		       __FUNCTION__, __LINE__);
		err = 0;
		goto free_key;
	}

	err = crypto_akcipher_set_pub_key(tfm, key, vec->key_len);
	if (err) {
		printk(KERN_ERR
		       "%s line %d failed to set crypto_akcipher_set_pub_key ",
		       __FUNCTION__, __LINE__);
		goto free_key;
	}

	/*
	 * First run test which do not require a private key, such as
	 * encrypt or verify.
	 */
	err = -ENOMEM;
	out_len_max = crypto_akcipher_maxsize(tfm);
	src_buffer = kzalloc(out_len_max, GFP_KERNEL);
	if (!src_buffer)
		goto free_key;

	err = -E2BIG;

	memcpy(src_buffer, vec->c, vec->c_size); // signature + E

	sg_init_table(src_tab, 3);
	sg_set_buf(&src_tab[0], src_buffer, 8);
	sg_set_buf(&src_tab[1], src_buffer + 8, vec->c_size - 8);

	{
		memcpy(src_buffer + vec->c_size, vec->m, vec->m_size);
		sg_set_buf(&src_tab[2], src_buffer + vec->c_size, vec->m_size);
		akcipher_request_set_crypt(req, src_tab, NULL, vec->c_size,
					   vec->m_size);
	}

	akcipher_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				      crypto_req_done, &wait);

	err = crypto_wait_req(crypto_akcipher_verify(req), &wait);
	if (err) {
		printk(KERN_ERR "alg: akcipher: verify test failed. err %d\n",
		       err);
		goto free_all;
	}

free_all:
	kfree(src_buffer);
	kfree(dst_buffer);
free_key:
	kfree(key);

free_req:
	akcipher_request_free(req);

	crypto_free_akcipher(tfm);

	return err;
}

static uint8_t sign[1024] = { 0 };
static uint8_t cipher[1024] = { 0 };
int sm2_test(void)
{
	int i;
	struct akcipher_testvec *test_vecs =
		(struct akcipher_testvec *)sm2_tv_template;
	struct akcipher_testvec *vec;
	u32 count = sizeof(sm2_tv_template) / sizeof(sm2_tv_template[0]);
	int ret;
	int err = 0;
	struct akcipher_testvec tmp_vec = { 0 };

	vec = &tmp_vec;
	for (i = 0; i < count; i++) {
		if (idx != 0xffff && i != idx) {
			pr_info("Skip sm2 [%u]\n", i);
			continue;
		}

		memcpy(&tmp_vec, &test_vecs[i], sizeof(tmp_vec));

		if (vec->public_key_vec) { // verify ?
			if (vec->siggen_sigver_test) {
				ret = ecdsa_test_verify("sm2", vec);
				if (ret) {
					pr_err("Test error [%d]\r\n", i);
					err++;
				} else {
					pr_info("Test verify success[%d]\r\n",
						i);
				}
			} else {
				vec->c = cipher;
				vec->c_size = 1024;
				ret = ecdsa_test_encrypt("sm2", vec);
				if (ret) {
					pr_err("Test encrypt error [%d]\r\n",
					       i);
					err++;
				}

				vec->key = vec->priv_key;
				vec->key_len = 32;
				vec->public_key_vec = false;
				ret = ecdsa_test_decrypt("sm2", vec);
				if (ret) {
					pr_err("Test decrypt cipher error [%d]\r\n",
					       i);
					err++;
				} else {
					pr_info("Test encrypt success[%d]\r\n",
						i);
				}
			}
		} else {
			if (vec->siggen_sigver_test) {
				vec->c = sign;
				ret = ecdsa_test_sign("sm2", vec);
				vec->public_key_vec = true;
				vec->key = vec->pub_key;
				vec->key_len = 65;
				ret = ecdsa_test_verify("sm2", vec);
				if (ret) {
					pr_err("Verify error for sign [%d]\r\n",
					       i);
					err++;
				} else {
					pr_info("Test sign success[%d]\r\n", i);
				}
			} else {
				ret = ecdsa_test_decrypt("sm2", vec);
				if (ret) {
					pr_err("Test decrypt error [%d]\r\n",
					       i);
					err++;
				} else {
					pr_info("Test decrypt success[%d]\r\n",
						i);
				}
			}
		}
	}

	return err;
}

void sm2_speed(int secs)
{
	unsigned long start, end;
	int count, ret, i;
	struct akcipher_testvec enc_vec;
	struct akcipher_testvec dec_vec;
	struct akcipher_testvec sign_vec;
	struct akcipher_testvec verify_vec;

	u32 tv_count = sizeof(sm2_tv_template) / sizeof(sm2_tv_template[0]);

	for (i = 0; i < tv_count; i++) {
		if (sm2_tv_template[i].public_key_vec) {
			if (sm2_tv_template[i].siggen_sigver_test) {
				memcpy(&verify_vec, &sm2_tv_template[i],
				       sizeof(struct akcipher_testvec));
			} else {
				memcpy(&enc_vec, &sm2_tv_template[i],
				       sizeof(struct akcipher_testvec));
			}
		} else {
			if (sm2_tv_template[i].siggen_sigver_test) {
				memcpy(&sign_vec, &sm2_tv_template[i],
				       sizeof(struct akcipher_testvec));

			} else {
				memcpy(&dec_vec, &sm2_tv_template[i],
				       sizeof(struct akcipher_testvec));
			}
		}
	}

	/*
	encrypt speed
	*/
	for (start = jiffies, end = start + secs * HZ, count = 0;
	     time_before(jiffies, end); count++) {
		enc_vec.c = cipher;
		enc_vec.c_size = 1024;
		ret = ecdsa_test_encrypt("sm2", &enc_vec);

		if (ret)
			return;
	}
	pr_info("Doing %s for %us : %6u opers in %d s %6u opers/sec\n",
		"SM2 encrypt", secs, count, secs, count / secs);

	/*
	decrypt speed
	*/
	for (start = jiffies, end = start + secs * HZ, count = 0;
	     time_before(jiffies, end); count++) {
		ret = ecdsa_test_decrypt("sm2", &dec_vec);

		if (ret)
			return;
	}
	pr_info("Doing %s for %us: %6u opers in %d s %6u opers/sec\n",
		"SM2 decrypt", secs, count, secs, count / secs);

	/*
	sign speed
	*/
	sign_vec.c = sign;
	for (start = jiffies, end = start + secs * HZ, count = 0;
	     time_before(jiffies, end); count++) {
		ret = ecdsa_test_sign("sm2", &sign_vec);

		if (ret)
			return;
	}
	pr_info("Doing %s for %us : %6u opers in %d s %6u opers/sec\n",
		"SM2 sign", secs, count, secs, count / secs);

	/*
	verify speed
	*/
	for (start = jiffies, end = start + secs * HZ, count = 0;
	     time_before(jiffies, end); count++) {
		ret = ecdsa_test_verify("sm2", &verify_vec);

		if (ret)
			return;
	}
	pr_info("Doing %s for %us: %6u opers in %d s %6u opers/sec\n",
		"SM2 verify", secs, count, secs, count / secs);
}
